import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict

from models.blocks.model.attention.SelfAttention import ScaledDotProductAttention
from models.blocks.model.attention.SimplifiedSelfAttention import SimplifiedScaledDotProductAttention

__all__ = ['ExternalAttention', 'SEAttention', 'ChannelSpatialAttentionModule', 'PartExternalAttention',
           'PartSpatialAttention', 'SKAttention', 'ECAAttention']


class PartExternalAttention(nn.Module):
    def __init__(self, spa_feature_length, batch_feature_length, spa_S, batch_S):
        super().__init__()
        self.spa_mk = nn.Linear(spa_feature_length, spa_S, bias=False)
        self.spa_mv = nn.Linear(spa_S, spa_feature_length, bias=False)
        self.softmax = nn.Softmax(dim=1)
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, feature_map):
        b_s, map_size, h, w = feature_map.shape
        # batch attention
        spatial_feature = feature_map.view(b_s, map_size, h * w).permute(0, 2, 1)
        spatial_feature = self.spatialAttention(spatial_feature)
        spatial_feature = spatial_feature.permute(0, 1, 2).view(b_s, map_size, h, w)
        return feature_map + spatial_feature

    def spatialAttention(self, queries):
        attn = self.spa_mk(queries)  # bs,n,S
        attn = self.softmax(attn)  # bs,n,S
        attn = attn / torch.sum(attn, dim=2, keepdim=True)  # bs,n,S
        out = self.spa_mv(attn)  # bs,n,d_model
        return out


class PartSpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], dim=1)
        output = self.conv(result)
        output = self.sigmoid(output)
        out = output * x
        return x + out


class ExternalAttention(nn.Module):

    def __init__(self, d_model, S=64):
        super().__init__()
        self.mk = nn.Linear(d_model, S, bias=False)
        self.mv = nn.Linear(S, d_model, bias=False)
        self.softmax = nn.Softmax(dim=1)
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, queries):
        bs, n, h, w = queries.shape
        queries = queries.view(bs, n, h * w).permute(0, 2, 1)
        attn = self.mk(queries)  # bs,n,S
        attn = self.softmax(attn)  # bs,n,S
        attn = attn / torch.sum(attn, dim=2, keepdim=True)  # bs,n,S
        out = self.mv(attn)  # bs,n,d_model
        out = out.permute(0, 2, 1).view(bs, n, h, w)
        return out


class SEAttention(nn.Module):

    def __init__(self, channel=512, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x + x * y.expand_as(x)


class ChannelSpatialAttentionModule(nn.Module):
    class ChannelAttention(nn.Module):
        def __init__(self, channel, reduction=16):
            super().__init__()
            self.maxpool = nn.AdaptiveMaxPool2d(1)
            self.avgpool = nn.AdaptiveAvgPool2d(1)
            self.se = nn.Sequential(
                nn.Conv2d(channel, channel // reduction, 1, bias=False),
                nn.ReLU(),
                nn.Conv2d(channel // reduction, channel, 1, bias=False)
            )
            self.sigmoid = nn.Sigmoid()

        def forward(self, x):
            max_result = self.maxpool(x)
            avg_result = self.avgpool(x)
            max_out = self.se(max_result)
            avg_out = self.se(avg_result)
            output = self.sigmoid(max_out + avg_out)
            return output

    class SpatialAttention(nn.Module):
        def __init__(self, kernel_size=7):
            super().__init__()
            self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
            self.sigmoid = nn.Sigmoid()

        def forward(self, x):
            max_result, _ = torch.max(x, dim=1, keepdim=True)
            avg_result = torch.mean(x, dim=1, keepdim=True)
            result = torch.cat([max_result, avg_result], dim=1)
            output = self.conv(result)
            output = self.sigmoid(output)
            return output

    def __init__(self, channel, reduction=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelSpatialAttentionModule.ChannelAttention(channel=channel, reduction=reduction)
        self.sa = ChannelSpatialAttentionModule.SpatialAttention(kernel_size=kernel_size)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        out = x * self.ca(x)
        out = out * self.sa(out)
        return x + out


class SKAttention(nn.Module):

    def __init__(self, channel=512, kernels=None, reduction=16, group=1, L=32):
        super().__init__()
        if kernels is None:
            kernels = [1, 3, 5, 7]
        self.d = max(L, channel // reduction)
        self.convs = nn.ModuleList([])
        for k in kernels:
            self.convs.append(
                nn.Sequential(OrderedDict([
                    ('conv', nn.Conv2d(channel, channel, kernel_size=k, padding=k // 2, groups=group)),
                    ('bn', nn.BatchNorm2d(channel)),
                    ('relu', nn.ReLU())
                ]))
            )
        self.fc = nn.Linear(channel, self.d)
        self.fcs = nn.ModuleList([])
        for i in range(len(kernels)):
            self.fcs.append(nn.Linear(self.d, channel))
        self.softmax = nn.Softmax(dim=0)

    def forward(self, x):
        bs, c, _, _ = x.size()
        conv_outs = []
        ### split
        for conv in self.convs:
            conv_outs.append(conv(x))
        feats = torch.stack(conv_outs, 0)  # k,bs,channel,h,w

        ### fuse
        U = sum(conv_outs)  # bs,c,h,w

        ### reduction channel
        S = U.mean(-1).mean(-1)  # bs,c
        Z = self.fc(S)  # bs,d

        ### calculate attention weight
        weights = []
        for fc in self.fcs:
            weight = fc(Z)
            weights.append(weight.view(bs, c, 1, 1))  # bs,channel
        attention_weughts = torch.stack(weights, 0)  # k,bs,channel,1,1
        attention_weughts = self.softmax(attention_weughts)  # k,bs,channel,1,1

        ### fuse
        V = (attention_weughts * feats).sum(0)
        return V


class ECAAttention(nn.Module):

    def __init__(self, kernel_size=3):
        super().__init__()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2)
        self.sigmoid = nn.Sigmoid()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        y = self.gap(x)  # bs,c,1,1
        y = y.squeeze(-1).permute(0, 2, 1)  # bs,1,c
        y = self.conv(y)  # bs,1,c
        y = self.sigmoid(y)  # bs,1,c
        y = y.permute(0, 2, 1).unsqueeze(-1)  # bs,c,1,1
        return x * y.expand_as(x)


class DAModule(nn.Module):
    class PositionAttentionModule(nn.Module):

        def __init__(self, d_model=512, kernel_size=3, H=7, W=7):
            super().__init__()
            self.cnn = nn.Conv2d(d_model, d_model, kernel_size=kernel_size, padding=(kernel_size - 1) // 2)
            self.pa = ScaledDotProductAttention(d_model, d_k=d_model, d_v=d_model, h=1)

        def forward(self, x):
            bs, c, h, w = x.shape
            y = self.cnn(x)
            y = y.view(bs, c, -1).permute(0, 2, 1)  # bs,h*w,c
            y = self.pa(y, y, y)  # bs,h*w,c
            return y

    class ChannelAttentionModule(nn.Module):

        def __init__(self, d_model=512, kernel_size=3, H=7, W=7):
            super().__init__()
            self.cnn = nn.Conv2d(d_model, d_model, kernel_size=kernel_size, padding=(kernel_size - 1) // 2)
            self.pa = SimplifiedScaledDotProductAttention(H * W, h=1)

        def forward(self, x):
            bs, c, h, w = x.shape
            y = self.cnn(x)
            y = y.view(bs, c, -1)  # bs,c,h*w
            y = self.pa(y, y, y)  # bs,c,h*w
            return y

    def __init__(self, d_model=512, kernel_size=3, H=7, W=7):
        super().__init__()
        self.position_attention_module = DAModule.PositionAttentionModule(d_model, kernel_size, H, W)
        self.channel_attention_module = DAModule.ChannelAttentionModule(d_model, kernel_size, H, W)

    def forward(self, input):
        bs, c, h, w = input.shape
        p_out = self.position_attention_module(input)
        c_out = self.channel_attention_module(input)
        p_out = p_out.permute(0, 2, 1).view(bs, c, h, w)
        c_out = c_out.view(bs, c, h, w)
        return p_out + c_out
